/*
perone_avl - An AVL Search Tree algorithm jitter
Copyright (C) 2009 Christian S. Perone <christian.perone@gmail.com>
g_hash_table_new_full
This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License along
with this program; if not, write to the Free Software Foundation, Inc.,
51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
*/


/*
 * This program JITs an AVL Tree to perform faster
 * lookups. It uses a fixed function prototype and
 * it's just a PoC for JITting search alg. and 
 * their data structures.
*/


#include "perone_avl.h"

// Parameters
// Number of Nodes
#define TREE_SIZE 8
// Number of Lookups
#define LOOKUPS 1000
// Random Seed (for GRand)
#define SEED 54783

// Create a new module
LLVMModuleRef create_module(char *name)
{
   LLVMModuleRef mod = LLVMModuleCreateWithName(name);
   return mod;
}

// Create the function (using a fixed prototype)
LLVMValueRef create_function(LLVMModuleRef mod, char *name)
{
   LLVMTypeRef f_args[] = { LLVMInt32Type() };
   LLVMValueRef fac = LLVMAddFunction(mod, name, LLVMFunctionType(LLVMInt32Type(), f_args, 1, 0));
   LLVMSetFunctionCallConv(fac, LLVMCCallConv /*LLVMFastCallConv*/);
   LLVMSetLinkage(fac, LLVMInternalLinkage);
   return fac;
}

// Verify the new created module, abort on problems
void verify_module(LLVMModuleRef mod)
{
   char *error = NULL;
   LLVMVerifyModule(mod, LLVMAbortProcessAction, &error);
   LLVMDisposeMessage(error);
   return;
}

// This is the main function, it makes the codegen
// of the AVL Tree
LLVMModuleRef translate_avl_tree(GTree *t)
{
   LLVMTypeRef type_int = LLVMInt32Type();

   LLVMModuleRef mod       = create_module("avlmodule");
   LLVMValueRef func       = create_function(mod, "avllookup");
   LLVMValueRef lookup_arg = LLVMGetParam(func, 0);
   LLVMBasicBlockRef entry = LLVMAppendBasicBlock(func, "entry");

   LLVMBasicBlockRef basic_block       = NULL;
   LLVMBasicBlockRef basic_block_true  = NULL;
   LLVMBasicBlockRef basic_block_false = NULL;

   LLVMValueRef con = NULL;
   LLVMValueRef ret = NULL;
   LLVMValueRef ifx = NULL;

   LLVMBuilderRef builder = LLVMCreateBuilder();

   GList *stack           = NULL;
   GList *pre_order_stack = NULL;
   GList *list_item       = NULL;

   GTreeNode *node       = NULL;
   GTreeNode *node_left  = NULL;
   GTreeNode *node_right = NULL;

   GHashTable *label_mappings = g_hash_table_new_full(g_str_hash, g_str_equal, g_free, NULL);
   int node_left_key, node_right_key;
   char *buffer  = NULL;
   int pop_key   = 0;
   int pop_value = 0;

   buffer = g_strdup("BRNULL");
   basic_block = LLVMAppendBasicBlock(func, buffer);
   g_hash_table_insert(label_mappings, buffer, basic_block);
   LLVMPositionBuilderAtEnd(builder, basic_block);
   con = LLVMConstInt(type_int, -1, 0);
   LLVMBuildRet(builder, con);

   stack = g_list_append(stack, t->root);

   while(g_list_length(stack)>0)
   {
      node = (GTreeNode*) g_list_last(stack)->data;
      stack = g_list_remove(stack, node);
      pre_order_stack = g_list_append(pre_order_stack, node);

      if(node->right_child)
         stack = g_list_append(stack, node->right);
      if(node->left_child)
         stack = g_list_append(stack, node->left);
   }
   

   while((list_item = g_list_last(pre_order_stack)))
   {
      node = (GTreeNode*) list_item->data;
      pre_order_stack = g_list_remove(pre_order_stack, node);
      pop_key = GPOINTER_TO_INT(node->key);
      pop_value = GPOINTER_TO_INT(node->value);

      // Always create the EQX label ---------------------------------
      buffer = g_strdup_printf("EQ%d", pop_key);
      basic_block = LLVMAppendBasicBlock(func, buffer);
      g_hash_table_insert(label_mappings, buffer, basic_block);
      LLVMPositionBuilderAtEnd(builder, basic_block);
      con = LLVMConstInt(type_int, pop_value, 0);
      ret = LLVMBuildRet(builder, con);

      // Create the DIFX label -----------------------------------------
      buffer = g_strdup_printf("DIF%d", pop_key);
      basic_block = LLVMAppendBasicBlock(func, buffer);
      g_hash_table_insert(label_mappings, buffer, basic_block);
      LLVMPositionBuilderAtEnd(builder, basic_block);

      // 1) It's a pretty green leaf
      if((!node->right_child) && (!node->left_child))
      {
         // Get the BRNULL
         buffer = g_strdup("BRNULL");
         basic_block_false = (LLVMBasicBlockRef)g_hash_table_lookup(label_mappings, buffer);
         assert(basic_block_false != NULL);
         g_free(buffer);
         LLVMBuildBr(builder, basic_block_false);
      }
      else // It's an internal node
      {
         // We have a "greater than" node
         if(node->right_child)
         {
            node_right = node->right;
            node_right_key = GPOINTER_TO_INT(node_right->key);
            buffer = g_strdup_printf("BR%d", node_right_key);
            basic_block_true = (LLVMBasicBlockRef) g_hash_table_lookup(label_mappings, buffer);
            assert(basic_block_true != NULL);
            g_free(buffer);
         }
         else  // We do not have a "greater than" node
         {
            // Get the BRNULL
            buffer = g_strdup("BRNULL");
            basic_block_true = (LLVMBasicBlockRef) g_hash_table_lookup(label_mappings, buffer);
            assert(basic_block_true != NULL);
            g_free(buffer);
         }

         // We have a "less than" node
         if(node->left_child)
         {
            node_left = node->left;
            node_left_key = GPOINTER_TO_INT(node_left->key);
            buffer = g_strdup_printf("BR%d", node_left_key);
            basic_block_false = (LLVMBasicBlockRef) g_hash_table_lookup(label_mappings, buffer);
            assert(basic_block_false != NULL);
            g_free(buffer);
         } 
         else  // We do not have a "less than" node
         {
            // Get the BRNULL
            buffer = g_strdup("BRNULL");
            basic_block_false = (LLVMBasicBlockRef) g_hash_table_lookup(label_mappings, buffer);
            assert(basic_block_false != NULL);
            g_free(buffer);
         }

         ifx = LLVMBuildICmp(builder, LLVMIntSGT, lookup_arg,
                             LLVMConstInt(LLVMInt32Type(), pop_key, 0),
                             "Equality");

         LLVMBuildCondBr(builder, ifx, basic_block_true, basic_block_false);
      }

      // Create the BRX label -----------------------------------------
      // 1) It's the root node, so it's the "entry" label
      if(g_list_length(pre_order_stack)==0) basic_block = entry;
      else
      {
         buffer = g_strdup_printf("BR%d", pop_key);
         basic_block = LLVMAppendBasicBlock(func, buffer);
         g_hash_table_insert(label_mappings, buffer, basic_block);
      }

      LLVMPositionBuilderAtEnd(builder, basic_block);

      buffer = g_strdup_printf("EQ%d", pop_key);
      basic_block_true = (LLVMBasicBlockRef) g_hash_table_lookup(label_mappings, buffer);
      assert(basic_block_true != NULL);
      g_free(buffer);

      buffer = g_strdup_printf("DIF%d", pop_key);
      basic_block_false = (LLVMBasicBlockRef) g_hash_table_lookup(label_mappings, buffer);
      assert(basic_block_false != NULL);
      g_free(buffer);

      ifx = LLVMBuildICmp(builder, LLVMIntEQ, lookup_arg,
                          LLVMConstInt(LLVMInt32Type(), pop_key, 0),
                          "Equality");

      LLVMBuildCondBr(builder, ifx, basic_block_true, basic_block_false);

   }
   LLVMDisposeBuilder(builder);
   g_hash_table_destroy(label_mappings);
   g_list_free(stack);
   g_list_free(pre_order_stack);
   return mod;
}

void run_passes(LLVMModuleProviderRef mp,
                LLVMValueRef func,
                LLVMExecutionEngineRef engine)
{
   int error_ret = 0;
   LLVMPassManagerRef pm = NULL;

   pm = LLVMCreateFunctionPassManager(mp);
   LLVMAddTargetData(LLVMGetExecutionEngineTargetData(engine), pm);

   LLVMAddSCCPPass(pm);
   LLVMAddInstructionCombiningPass(pm);
   LLVMAddConstantPropagationPass(pm);
   LLVMAddCFGSimplificationPass(pm);
   LLVMAddScalarReplAggregatesPass(pm);
   LLVMAddSimplifyLibCallsPass(pm);
   LLVMAddInstructionCombiningPass(pm);
   LLVMAddJumpThreadingPass(pm);
   LLVMAddCFGSimplificationPass(pm);
   LLVMAddInstructionCombiningPass(pm);
   LLVMAddTailCallEliminationPass(pm);
   LLVMAddCFGSimplificationPass(pm);
   LLVMAddReassociatePass(pm);
   LLVMAddLoopRotatePass(pm);
   LLVMAddLICMPass(pm);
   LLVMAddLoopUnswitchPass(pm);
   LLVMAddInstructionCombiningPass(pm);
   LLVMAddIndVarSimplifyPass(pm);
   LLVMAddLoopDeletionPass(pm);
   LLVMAddLoopUnrollPass(pm);
   LLVMAddInstructionCombiningPass(pm);
   LLVMAddGVNPass(pm);
   LLVMAddMemCpyOptPass(pm);
   LLVMAddSCCPPass(pm);
   LLVMAddInstructionCombiningPass(pm);
   LLVMAddJumpThreadingPass(pm);
   LLVMAddDeadStoreEliminationPass(pm);
   LLVMAddAggressiveDCEPass(pm);
   LLVMAddCFGSimplificationPass(pm);
   LLVMAddPromoteMemoryToRegisterPass(pm);

   printf("\nInitialize Function Pass: ");
   error_ret = LLVMInitializeFunctionPassManager(pm);
   if(error_ret==1) printf("module modified !");
   else printf("module original !");

   printf("\nRun Function Pass: ");
   error_ret = LLVMRunFunctionPassManager(pm, func);
   if(error_ret==1) printf("module modified !");
   else printf("module original !");

   printf("\nFinalize Function Pass: ");
   error_ret = LLVMFinalizeFunctionPassManager(pm);
   if(error_ret==1) printf("module modified !");
   else printf("module original !");

   LLVMDisposePassManager(pm);
   return;
}


void get_execution_engine(LLVMModuleRef mod,
                          LLVMExecutionEngineRef *engine,
                          LLVMModuleProviderRef *provider)
{
   int error_ret;
   char *verror = NULL;
   *provider = LLVMCreateModuleProviderForExistingModule(mod);
   assert(provider!=NULL);

   error_ret = LLVMCreateJITCompiler(engine, *provider, 3, &verror);
   if(error_ret)
   {
      fprintf(stderr, "\nError code [%d] in JIT: [%s]\n", error_ret, verror);
      LLVMDisposeMessage(verror);
      abort();
   }
   return;
}

void *get_function_pointer(LLVMExecutionEngineRef engine,
                           LLVMValueRef func)
{
   void *func_ptr = NULL;
   func_ptr = LLVMGetPointerToGlobal(engine, func);
   return func_ptr;
}

double run_jit(LLVMExecutionEngineRef engine,
               void *func_ptr)
{

   jit_avl_lookup_t jit_avl_lookup = (jit_avl_lookup_t) func_ptr;

   int error_ret = 0;
   int i = 0;
   clock_t t0, t1;
   //int ran;
   double elapsed = 0.0;

   GRand *grand = g_rand_new_with_seed(SEED);

   printf("\nDoing %d lookups...", LOOKUPS);
   t0 = clock();
   for(i=0; i<LOOKUPS; i++)
   {
      //ran = g_rand_int_range(grand, 0, TREE_SIZE); 
      error_ret = jit_avl_lookup(g_rand_int_range(grand, 0, TREE_SIZE));
      //if(ran!=error_ret) printf("\nIntegrity error - %d / %d", ran, error_ret);
   }

   t1 = clock();
   elapsed = ((double) t1-t0)/CLOCKS_PER_SEC;
   printf("\nTook %.3f seconds !\n\n", elapsed);

   g_rand_free(grand);
   return elapsed;
}

int compare_int(const void *a, const void *b)
{
   return GPOINTER_TO_INT(a)-GPOINTER_TO_INT(b);
}

int main(int argc, char **argv)
{
   GTree *avl_tree = NULL;
   int i;
   void *lookup;
   void *func_ptr;
   clock_t t0, t1;
   GRand *grand;
   LLVMModuleRef mod;
   double nonjit_elapsed = 0.0;
   double jit_elapsed = 0.0;
   double jit_passes_elapsed = 0.0;
   double jit_compile_elapsed = 0.0;
   double total_jit = 0.0;
   
   LLVMExecutionEngineRef engine;
   LLVMModuleProviderRef provider;
   LLVMValueRef func;

   // Initialize the JIT
   LLVMLinkInJIT();
   LLVMInitializeNativeTarget();

   grand = g_rand_new_with_seed(SEED);

   printf("\nCreating tree of size %d...\n", TREE_SIZE);
   avl_tree = g_tree_new_full((GCompareDataFunc) compare_int,
                               NULL, NULL, NULL);

   for(i=0; i<TREE_SIZE; i++)
      g_tree_insert(avl_tree, GINT_TO_POINTER(i), GINT_TO_POINTER(i));

   printf("Tree info:\n");
   printf("Tree Height = %d\n", g_tree_height(avl_tree));
   printf("Tree Nodes  = %d\n", g_tree_nnodes(avl_tree));

   printf("\nDoing %d lookups...", LOOKUPS);
   t0 = clock();
   for(i=0; i<LOOKUPS; i++) 
      lookup = g_tree_lookup(avl_tree,
                             GINT_TO_POINTER(g_rand_int_range(grand, 0, TREE_SIZE)));

   t1 = clock();
   nonjit_elapsed = ((double) t1-t0)/CLOCKS_PER_SEC;
   printf("\nTook %.3f seconds !\n", nonjit_elapsed);

   printf("\nTranslating tree...");
   t0 = clock();
   mod = translate_avl_tree(avl_tree);
   t1 = clock();
   printf("\nTook %.3f seconds !\n", ((double) t1-t0)/CLOCKS_PER_SEC);

   LLVMWriteBitcodeToFile(mod, "outx.bc");

   get_execution_engine(mod, &engine, &provider);
   func = LLVMGetNamedFunction(mod, "avllookup");

   printf("\nRunning transformations...");
   t0 = clock();
   run_passes(provider, func, engine);
   t1 = clock();
   jit_passes_elapsed = ((double) t1-t0)/CLOCKS_PER_SEC;
   printf("\nTook %.3f seconds !\n\n", jit_passes_elapsed);

   printf("\nGetting function pointer...");
   t0 = clock();
   func_ptr = get_function_pointer(engine, func);
   t1 = clock();
   jit_compile_elapsed = ((double) t1-t0)/CLOCKS_PER_SEC;
   printf("\nTook %.3f seconds !\n\n", jit_compile_elapsed);
   
   jit_elapsed = run_jit(engine, func_ptr);

   total_jit = jit_elapsed+jit_compile_elapsed+jit_passes_elapsed;

   printf("\n--- Summary ---\n");
   printf("\nNon-JIT Time:                    %.3f", nonjit_elapsed);
   printf("\nJIT Time:                        %.3f", jit_elapsed);
   printf("\nJIT Compile Time:                %.3f", jit_compile_elapsed);
   printf("\nJIT Pass Time:                   %.3f", jit_passes_elapsed);
   printf("\nJIT + Pass + Compile:            %.3f", total_jit);
   printf("\nPerformance of JIT over non-JIT: %.3f%%", ((nonjit_elapsed-jit_elapsed)*100.0)/nonjit_elapsed);
   printf("\nPerf. of Total JIT over non-JIT: %.3f%%", ((nonjit_elapsed-total_jit)*100.0)/nonjit_elapsed);
   printf("\n\n");

   LLVMDisposeExecutionEngine(engine);
   g_tree_destroy(avl_tree);
   g_rand_free(grand);
   return 0;
}


